# coding:UTF-8
# IMU数据记录程序
import serial
import time
import csv
import os
from datetime import datetime

# --- IMU串口参数定义 ---
imu_params = {
    'port': '/dev/ttyUSB0',  # IMU串口号
    'baudrate': 9600,
    'timeout': 0.5
}

# --- 全局变量 ---
# IMU数据相关
ACCData = [0.0]*8
GYROData = [0.0]*8
AngleData = [0.0]*8
FrameState = 0  # 帧状态判断
Bytenum = 0  # 当前帧已接收字节数
CheckSum = 0  # 校验和

a = [0.0]*3  # 加速度
w = [0.0]*3  # 角速度
Angle = [0.0]*3  # 角度

# --- IMU数据处理函数 ---
def DueData(inputdata):
    """处理IMU原始数据
    Args:
        inputdata: IMU发送的原始字节数据
    Returns:
        bool: 是否成功解析出一组完整数据
    """
    global FrameState
    global Bytenum
    global CheckSum
    global a
    global w
    global Angle
    
    for data in inputdata:  # 遍历输入数据
        if FrameState == 0:  # 状态0：寻找帧头
            if data == 0x55 and Bytenum == 0:  # 找到第一个帧头字节0x55
                CheckSum = data
                Bytenum = 1
                continue
            elif data == 0x51 and Bytenum == 1:  # 加速度数据帧
                CheckSum += data
                FrameState = 1
                Bytenum = 2
            elif data == 0x52 and Bytenum == 1:  # 角速度数据帧
                CheckSum += data
                FrameState = 2
                Bytenum = 2
            elif data == 0x53 and Bytenum == 1:  # 角度数据帧
                CheckSum += data
                FrameState = 3
                Bytenum = 2
        elif FrameState == 1:  # 状态1：接收加速度数据
            if Bytenum < 10:  # 读取8个数据字节
                ACCData[Bytenum-2] = data
                CheckSum += data
                Bytenum += 1
            else:
                if data == (CheckSum & 0xff):  # 校验和检查
                    a = get_acc(ACCData)
                CheckSum = 0
                Bytenum = 0
                FrameState = 0
        elif FrameState == 2:  # 状态2：接收角速度数据
            if Bytenum < 10:
                GYROData[Bytenum-2] = data
                CheckSum += data
                Bytenum += 1
            else:
                if data == (CheckSum & 0xff):
                    w = get_gyro(GYROData)
                CheckSum = 0
                Bytenum = 0
                FrameState = 0
        elif FrameState == 3:  # 状态3：接收角度数据
            if Bytenum < 10:
                AngleData[Bytenum-2] = data
                CheckSum += data
                Bytenum += 1
            else:
                if data == (CheckSum & 0xff):
                    Angle = get_angle(AngleData)
                    return True  # 返回True表示成功解析了一组完整数据
                CheckSum = 0
                Bytenum = 0
                FrameState = 0
    return False  # 返回False表示没有解析到完整数据

def get_acc(datahex):
    """计算加速度值
    Args:
        datahex: 加速度的原始数据字节
    Returns:
        tuple: (x, y, z)轴的加速度
    """
    axl = datahex[0]
    axh = datahex[1]
    ayl = datahex[2]
    ayh = datahex[3]
    azl = datahex[4]
    azh = datahex[5]
    k_acc = 16.0
    acc_x = (axh << 8 | axl) / 32768.0 * k_acc
    acc_y = (ayh << 8 | ayl) / 32768.0 * k_acc
    acc_z = (azh << 8 | azl) / 32768.0 * k_acc
    if acc_x >= k_acc:
        acc_x -= 2 * k_acc
    if acc_y >= k_acc:
        acc_y -= 2 * k_acc
    if acc_z >= k_acc:
        acc_z -= 2 * k_acc
    return acc_x, acc_y, acc_z

def get_gyro(datahex):
    """计算角速度值
    Args:
        datahex: 角速度的原始数据字节
    Returns:
        tuple: (x, y, z)轴的角速度
    """
    wxl = datahex[0]
    wxh = datahex[1]
    wyl = datahex[2]
    wyh = datahex[3]
    wzl = datahex[4]
    wzh = datahex[5]
    k_gyro = 2000.0
    gyro_x = (wxh << 8 | wxl) / 32768.0 * k_gyro
    gyro_y = (wyh << 8 | wyl) / 32768.0 * k_gyro
    gyro_z = (wzh << 8 | wzl) / 32768.0 * k_gyro
    if gyro_x >= k_gyro:
        gyro_x -= 2 * k_gyro
    if gyro_y >= k_gyro:
        gyro_y -= 2 * k_gyro
    if gyro_z >= k_gyro:
        gyro_z -= 2 * k_gyro
    return gyro_x, gyro_y, gyro_z

def get_angle(datahex):
    """计算角度值
    Args:
        datahex: 角度的原始数据字节
    Returns:
        tuple: (x, y, z)轴的角度
    """
    rxl = datahex[0]
    rxh = datahex[1]
    ryl = datahex[2]
    ryh = datahex[3]
    rzl = datahex[4]
    rzh = datahex[5]
    k_angle = 180.0
    angle_x = (rxh << 8 | rxl) / 32768.0 * k_angle
    angle_y = (ryh << 8 | ryl) / 32768.0 * k_angle
    angle_z = (rzh << 8 | rzl) / 32768.0 * k_angle
    if angle_x >= k_angle:
        angle_x -= 2 * k_angle
    if angle_y >= k_angle:
        angle_y -= 2 * k_angle
    if angle_z >= k_angle:
        angle_z -= 2 * k_angle
    return angle_x, angle_y, angle_z

# --- 主程序 ---
if __name__ == '__main__':
    try:
        # 打开IMU串口
        ser_imu = serial.Serial(**imu_params)
        print(f"IMU串口已打开: {ser_imu.is_open}, 端口: {imu_params['port']}")
        
        # 创建CSV文件
        filename_input = input("请输入保存文件的名称(如输入数字则保存为'imu_数字.csv'，输入'q'退出): ")
        if filename_input.lower() == 'q':
            print("程序退出")
            exit(0)
            
        filename = f"imu_{filename_input}.csv"
        print(f"将保存数据到文件: {filename}")
        
        with open(filename, 'w', newline='') as csvfile:
            csv_writer = csv.writer(csvfile)
            csv_writer.writerow([
                'Timestamp', 
                'IMU_Acc_X', 'IMU_Acc_Y', 'IMU_Acc_Z', 
                'IMU_Gyro_X', 'IMU_Gyro_Y', 'IMU_Gyro_Z',
                'IMU_Angle_X', 'IMU_Angle_Y', 'IMU_Angle_Z'
            ])
            
            # 设置运行时间
            run_time = 5.0  # 5秒
            start_time = time.time()
            data_count = 0
            
            print(f"开始记录数据，将运行{run_time}秒...")
            
            # 主循环
            while time.time() - start_time < run_time:
                # 读取IMU数据
                datahex = ser_imu.read(33)
                
                # 处理IMU数据
                if datahex and DueData(datahex):
                    # 显示IMU数据
                    print("加速度: %10.3f %10.3f %10.3f" % a)
                    print("角速度: %10.3f %10.3f %10.3f" % w)
                    print("角度: %10.3f %10.3f %10.3f" % Angle)
                    print("-" * 50)
                    
                    # 记录数据
                    timestamp = time.time() - start_time
                    csv_writer.writerow([
                        timestamp,
                        a[0], a[1], a[2],
                        w[0], w[1], w[2],
                        Angle[0], Angle[1], Angle[2]
                    ])
                    csvfile.flush()
                    data_count += 1
                    
                    if data_count % 10 == 0:
                        elapsed = time.time() - start_time
                        remaining = max(0, run_time - elapsed)
                        print(f"已记录 {data_count} 组数据, 剩余时间: {remaining:.1f}秒")
                
                # 短暂休眠以减少CPU使用率
                time.sleep(0.01)
            
            print(f"共记录 {data_count} 组数据")
            
            # 添加等待用户输入，让用户选择继续记录或退出
            while True:
                next_action = input("记录完成。输入新的文件名继续记录，或输入'q'退出: ")
                if next_action.lower() == 'q':
                    break
                    
                # 新文件名
                filename = f"imu_{next_action}.csv"
                print(f"将保存数据到新文件: {filename}")
                
                with open(filename, 'w', newline='') as csvfile:
                    csv_writer = csv.writer(csvfile)
                    csv_writer.writerow([
                        'Timestamp', 
                        'IMU_Acc_X', 'IMU_Acc_Y', 'IMU_Acc_Z', 
                        'IMU_Gyro_X', 'IMU_Gyro_Y', 'IMU_Gyro_Z',
                        'IMU_Angle_X', 'IMU_Angle_Y', 'IMU_Angle_Z'
                    ])
                    
                    # 重置计数和时间
                    start_time = time.time()
                    data_count = 0
                    
                    print(f"开始记录新数据，将运行{run_time}秒...")
                    
                    # 主循环
                    while time.time() - start_time < run_time:
                        # 读取IMU数据
                        datahex = ser_imu.read(33)
                        
                        # 处理IMU数据
                        if datahex and DueData(datahex):
                            # 显示IMU数据
                            print("加速度: %10.3f %10.3f %10.3f" % a)
                            print("角速度: %10.3f %10.3f %10.3f" % w)
                            print("角度: %10.3f %10.3f %10.3f" % Angle)
                            print("-" * 50)
                            
                            # 记录数据
                            timestamp = time.time() - start_time
                            csv_writer.writerow([
                                timestamp,
                                a[0], a[1], a[2],
                                w[0], w[1], w[2],
                                Angle[0], Angle[1], Angle[2]
                            ])
                            csvfile.flush()
                            data_count += 1
                            
                            if data_count % 10 == 0:
                                elapsed = time.time() - start_time
                                remaining = max(0, run_time - elapsed)
                                print(f"已记录 {data_count} 组数据, 剩余时间: {remaining:.1f}秒")
                        
                        # 短暂休眠以减少CPU使用率
                        time.sleep(0.01)
                    
                    print(f"共记录 {data_count} 组数据")
    
    except KeyboardInterrupt:
        print("\n程序被用户中断")
    
    except serial.SerialException as e:
        print(f"串口错误: {e}")
    
    except Exception as e:
        print(f"程序错误: {e}")
    
    finally:
        # 关闭串口
        if 'ser_imu' in locals() and ser_imu.is_open:
            ser_imu.close()
            print("IMU串口已关闭") 